"""
BEL 标签器

使用二元错误定位算法为推理步骤添加标签，严格按照设计文档要求
复用现有的 vLLM 推理框架和统一的工具函数
"""

import json
import os
import sys
from typing import Dict, List, Optional, Any
from pathlib import Path
import time

# 复用现有的推理模块
from vllm import LLM, SamplingParams
from inference.vllm_infer import extract_answer_try_all_methods, template, prompt
from PIL import Image
import jsonlines
from loguru import logger

from .engine import BinaryErrorLocator
from .utils import (
    validate_bel_format, ensure_dir, count_lines, ProgressBar,
    format_steps_with_labels, split_steps, insert_labels,
    load_config, validate_config, setup_logging
)


class BinaryErrorLabeler:
    """
    推理步骤标签器
    
    使用二元错误定位算法为推理步骤添加标签
    严格按照设计文档要求，复用统一的工具函数
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化标签器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.error_locator = BinaryErrorLocator(config)
        
        # 从配置中提取参数
        self.batch_size = config['algorithm']['batch_size']
        
        logger.info(f"BinaryErrorLabeler 初始化完成")
    
    def label_rollouts(self, input_file: str, output_file: str, resume: bool = True) -> None:
        """
        为推理路径添加标签
        
        Args:
            input_file: 输入文件路径（rollouts.jsonl）
            output_file: 输出文件路径（labeled.jsonl）
            resume: 是否支持断点续跑
        """
        logger.info(f"开始标注推理路径: {input_file} -> {output_file}")
        
        # 确保输出目录存在
        ensure_dir(output_file)
        
        # 检查断点续跑
        start_idx = 0
        if resume and os.path.exists(output_file):
            start_idx = count_lines(output_file)
            logger.info(f"检测到已有输出文件，从第 {start_idx + 1} 条记录开始")
        
        # 加载数据
        data = self._load_rollouts(input_file)
        
        if start_idx >= len(data):
            logger.info("所有数据已处理完成")
            return
        
        # 过滤需要处理的数据
        data = data[start_idx:]
        
        logger.info(f"需要处理 {len(data)} 条数据")
        
        # 批量处理
        with ProgressBar(len(data), "标注推理步骤") as pbar:
            with open(output_file, 'a', encoding='utf-8') as f:
                for i in range(0, len(data), self.batch_size):
                    batch_data = data[i:i + self.batch_size]
                    
                    try:
                        # 处理批次
                        for record in batch_data:
                            # 验证数据格式
                            if not validate_bel_format(record):
                                logger.warning("跳过格式不正确的记录")
                                pbar.update(1)
                                continue
                            
                            # 使用错误定位器生成标签
                            all_labels = self.error_locator.locate_error(record)
                            
                            # 为每个rollout添加标签
                            labeled_rollouts = []
                            rollouts = record.get('rollouts', [])
                            
                            for j, (rollout, labels) in enumerate(zip(rollouts, all_labels)):
                                output_text = rollout.get('output', '')
                                
                                # 添加标签
                                if labels:
                                    labeled_output = insert_labels(output_text, labels)
                                else:
                                    labeled_output = output_text
                                
                                labeled_rollouts.append({
                                    'output': labeled_output,
                                    'labels': labels
                                })
                            
                            # 构建结果记录
                            result = {
                                **record,  # 保留原始数据
                                'rollouts': labeled_rollouts
                            }
                            
                            # 写入文件
                            f.write(json.dumps(result, ensure_ascii=False) + '\n')
                            f.flush()
                            
                            pbar.update(1)
                            
                    except Exception as e:
                        logger.error(f"批次 {i//self.batch_size + 1} 处理失败: {e}")
                        # 跳过失败的批次，避免中断整个流程
                        for _ in range(len(batch_data)):
                            pbar.update(1)
                        continue
        
        logger.info(f"推理步骤标注完成: {output_file}")
    
    def run(self):
        """
        运行标注器
        """
        output_config = self.config['output']
        
        input_file = output_config['rollouts']
        output_file = output_config['labeled']
        
        self.label_rollouts(input_file, output_file)
    
    def _load_rollouts(self, input_file: str) -> List[Dict]:
        """
        加载rollouts数据
        
        Args:
            input_file: 输入文件路径
            
        Returns:
            List[Dict]: rollouts数据列表
        """
        logger.info(f"加载rollouts数据: {input_file}")
        
        if not os.path.exists(input_file):
            raise FileNotFoundError(f"输入文件不存在: {input_file}")
        
        data = []
        try:
            if input_file.endswith('.json'):
                with open(input_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if not isinstance(data, list):
                        data = [data]  # 单个记录转换为列表
            elif input_file.endswith('.jsonl'):
                with open(input_file, 'r', encoding='utf-8') as f:
                    for line_num, line in enumerate(f, 1):
                        try:
                            item = json.loads(line.strip())
                            data.append(item)
                        except json.JSONDecodeError as e:
                            logger.warning(f"第 {line_num} 行 JSON 解析失败: {e}")
                            continue
            else:
                raise ValueError(f"不支持的文件格式: {input_file}")
                
        except Exception as e:
            logger.error(f"加载rollouts数据失败: {e}")
            raise
        
        logger.info(f"成功加载 {len(data)} 条rollouts数据")
        return data


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="BEL 二元错误标注器")
    parser.add_argument("--config", type=str, required=True, help="配置文件路径")
    parser.add_argument("--log_level", type=str, default="INFO", help="日志级别")
    
    args = parser.parse_args()
    
    # 设置日志
    setup_logging(args.log_level)
    
    # 加载配置
    config = load_config(args.config)
    if not validate_config(config):
        logger.error("配置文件验证失败")
        sys.exit(1)
    
    # 创建标注器并运行
    labeler = BinaryErrorLabeler(config)
    
    try:
        start_time = time.time()
        labeler.run()
        elapsed_time = time.time() - start_time
        logger.info(f"标注完成，耗时: {elapsed_time:.2f} 秒")
    except KeyboardInterrupt:
        logger.info("用户中断，程序退出")
    except Exception as e:
        logger.error(f"标注失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main() 